import os
import argparse

FILE_ENDWITH = "_HybrdChip"

CRC = '00110000000000000000000000000001'
CMD_RCRC_01 = '00110000000000001000000000000001'
CMD_RCRC_02 = '00000000000000000000000000000111'
COR1 = '00110000000000011100000000000001'

CMD_MASK_01 = '00110000000000001100000000000001'
CMD_MASK_02 = '10000000000000000000000000000000'
CMD_TRIM_01 = '00110000000000110110000000000001'
CMD_TRIM_02 = '10000000000000000000000000000000'

BITS_CMD = 'Bits:'

NOOP = '00100000000000000000000000000000'

def create_new_rbt_file(origin_file_path, new_file_path, force_write=True):        
    new_lines = []
    flag_COR1 = False
    flag_status = False # False is not modify / True is modify
    with open(origin_file_path, "r") as f:
        origin_lines = f.readlines()
        i = 0
        origin_lines_len = len(origin_lines)
        end_first_segment = 101
        while i < end_first_segment:
            line = origin_lines[i].strip()
            if CRC in line:
                if (i + 2) < end_first_segment and (i + 3) < end_first_segment:
                    if NOOP in origin_lines[i + 2] and NOOP in origin_lines[i + 3]:
                        new_lines.extend([CMD_RCRC_01, CMD_RCRC_02])
                        i += 2
                        continue
                new_lines.append(line)
                i += 1
                continue
            
            if COR1 in line and (i+1) < end_first_segment:
                flag_COR1 = True
                new_lines.append(line)
                i += 1
                continue

            if flag_COR1:
                flag_COR1 = False
                print(line)
                if line[-13] != "1":
                    flag_status = True
                    new_line = line[:-13] + "1" + line[-12:]
                    new_lines.extend([new_line, CMD_MASK_01, CMD_MASK_02, CMD_TRIM_01, CMD_TRIM_02])
                    i += 1
                else:
                    new_lines.append(line)
                    i += 1
            else:
                new_lines.append(line)
                i += 1

        # 处理中间部分
        start_second_segment = min(101, origin_lines_len)
        end_second_segment = max(0, origin_lines_len - 600)
        for i in range(start_second_segment, end_second_segment):
            line = origin_lines[i].strip()
            new_lines.append(line)

        # 处理后 600 行
        i = end_second_segment
        while i < origin_lines_len:
            line = origin_lines[i].strip()
            if CRC in line:
                if (i + 2) < origin_lines_len and (i + 3) < origin_lines_len:
                    if NOOP in origin_lines[i + 2] and NOOP in origin_lines[i + 3]:
                        new_lines.extend([CMD_RCRC_01, CMD_RCRC_02])
                        i += 2
                        continue
                new_lines.append(line)
                i += 1
                continue
            
            if COR1 in line and (i+1) < origin_lines_len:
                flag_COR1 = True
                new_lines.append(line)
                i += 1
                continue

            if flag_COR1:
                flag_COR1 = False
                if line[-13] != "1":
                    flag_status = True
                    new_line = line[:-13] + "1" + line[-12:]
                    new_lines.extend([new_line, CMD_MASK_01, CMD_MASK_02, CMD_TRIM_01, CMD_TRIM_02])
                    i += 1
                else:
                    new_lines.append(line)
                    i += 1
            else:
                new_lines.append(line)
                i += 1

    if force_write or flag_status:
        with open(new_file_path, 'w') as f:
            bits_flag = True
            for line in new_lines:
                if bits_flag and BITS_CMD in line:
                    bits_flag = False
                    old_bits_num = int(line.split(':')[1].strip())
                    new_bits_num = old_bits_num + 32 * (len(new_lines) - len(origin_lines))
                    line = f"Bits:\t{new_bits_num}"
                f.write(line + '\n')
        print(f"writing rbt file : {new_file_path}")

def find_rbt_files(root_path): 
    if not os.path.exists(root_path):
        print(f"Error: The specified path '{root_path}' does not exist.")
        return []
    # rsplit(".rbt", 1) right and only one
    rbt_files = [
        os.path.join(dirpath, filename).rsplit(".rbt", 1)[0]
        for dirpath, _, filenames in os.walk(root_path)
        for filename in filenames if filename.lower().endswith('.rbt')
    ]
    return rbt_files

def main():
    parser = argparse.ArgumentParser(description="Modify rbt trim")

    # Add optional arguments
    parser.add_argument('-r', '--rbt_folder', type=str, help="Input .rbt folder path")
    parser.add_argument('--force_write', action='store_true', default=True, help="Force write new file even if content is the same (default: True)")
    parser.add_argument('--no_force_write', action='store_false', dest='force_write', help="Do not force write new file even if content is the same")
    parser.add_argument('--only_one', type=str, help="Only process this specific .rbt file")
    parser.add_argument('--file_suffix', type=str, default=FILE_ENDWITH, help="Suffix to add to the new .rbt file (default: _HybrdChip)")

    # Parse arguments
    args = parser.parse_args()


    # Display help if no arguments are provided
    if not any(vars(args).values()):
        parser.print_help()
        return

    print(f"Parameters:")
    print(f"  rbt_folder: {args.rbt_folder}")
    print(f"  force_write: {args.force_write}")
    print(f"  only_one: {args.only_one}")
    print(f"  file_suffix: {args.file_suffix}\n")

    if args.only_one:
        # Process only the specified file
        origin_file_path = args.only_one
        new_file_path = origin_file_path + args.file_suffix + ".rbt"
        create_new_rbt_file(origin_file_path, new_file_path, force_write=args.force_write)
    if args.rbt_folder:
        # Process all files in the folder
        rbt_files = find_rbt_files(args.rbt_folder)
        for path in rbt_files:
            create_new_rbt_file(path+".rbt", path + args.file_suffix + ".rbt", force_write=args.force_write)
    if args.only_one == None and args.rbt_folder == None:
        print("No input folder or file specified. Use -h or --help for usage information.")

if __name__ == "__main__":
    main()
